from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.documents import Document
from langchain.chains.combine_documents import create_stuff_documents_chain
import os

os.environ['OPENAI_API_KEY'] = 'sk-ioXCJ3W2IQlpDzHu268aEeF7CcAf4b1f9f8c46A248C74a9d'
os.environ['OPENAI_BASE_URL'] = 'https://api.xiaoai.plus/v1'

if __name__ == '__main__':

    llm = ChatOpenAI()
    llm.model_name = 'gpt-4'

    prompt = ChatPromptTemplate.from_template(
        """仅根据提供的上下文回答以下问题:
        <上下文>
        {context}
        </上下文>
        问题: {input}"""
    )

    document_chain = create_stuff_documents_chain(llm,prompt)

    result = document_chain.invoke({
        "input": "重大哪一年被列入985工程？",
        "context": [Document(
            page_content="重庆大学（ChongqingUniversity，CQU），简称【重大】，是中华人民共和国教育部直属，由教育部、重庆市、国家国防科技工业局共建的全国重点大学，位列国家【双一流】、【211工程】、【985工程】，入选【珠峰计划】、【强基计划】、【高等学校创新能力提升计划】、【高等学校学科创新引智计划】、【卓越工程师教育培养计划】、【卓越法律人才教育培养计划】、国家建设高水平大学公派研究生项目、中国政府奖学金来华留学生接收院校、教育部来华留学示范基地，为卓越大学联盟、中波大学联盟、一带一路高校联盟、长江—伏尔加河”高校联盟、CDIO工程教育联盟、中国高等戏剧教育联盟成员单位、【国优计划】首批试点高校。学校创办于1929年，1942年成为一所综合性大学，1952年调整为多科性大学，1960年被确定为全国重点大学，1998年成为【211工程】重点建设高校；2000年，重庆大学合并重庆建筑大学、重庆建筑高等专科学校，2001年列入【985工程】重点建设高校，2017年入选【双一流】建设高校（A类）。截至2023年11月，学校有A校区、B校区、C校区、虎溪校区和两江校区，校园占地面积5300余亩；设有7个学部、35个学院及7所附属医院，开设108个本科专业；拥有35个博士后科研流动站，一级学科博士学位授权点38个、博士专业学位授权点8个，一级学科硕士学位授权点51个，专业硕士学位授权点25个；有教职工5300余人，其中专任教师3200余人；在校学生50000余人，其中研究生24000余人，本科生26000余人。")]
    })

    print(result)

